
#include "cumsum_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    CumsumTilingData tiling;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto coreNum = ascendcPlatform.GetCoreNum();
    auto BLOCK_DIM = 1;
    context->SetBlockDim(BLOCK_DIM);
    
    const auto inputDataType=context->GetInputTensor(0)->GetDataType();
    uint32_t typeSize=0;
    uint64_t key=1;
    //buffer for queue
	uint64_t UB_SHARING_NUM=4;
    switch(inputDataType){
        case ge::DT_FLOAT:
            typeSize=sizeof(float);
            key=1;
            break;
        case ge::DT_FLOAT16:
            typeSize=2;
            UB_SHARING_NUM=4;
            key=2;
	        break;
        case ge::DT_INT32:
            typeSize=sizeof(int32_t);
            UB_SHARING_NUM=7;
            key=3;
	        break; 
        case ge::DT_INT8:
            typeSize=sizeof(int8_t);
            UB_SHARING_NUM=7;
            key=4;
	        break; 
        default:
            //ASSERT(0 && "Unknown inputDataType");
            typeSize=4;
            break;
    }

    
    const auto inputShape=context->GetInputTensor(0)->GetOriginShape();
    //class Shape: size_t dim_num_; int64_t dims_[];
    int64_t dimSize=inputShape.GetDimNum();
    int64_t totalLength=context->GetInputTensor(0)->GetShapeSize();
    //axis in [-n,n-1]
    int32_t axis=*(context->GetInputTensor(1)->GetData<int32_t>());
    if(axis<0){
        axis+=dimSize;
    }
    //batchCnt for repeat times, batchLength means vector calculation length
    int64_t batchCnt=1;
    int64_t batchLength=1;
    int32_t i=0;
    for(i=0;i<axis;i++){
        batchCnt*=inputShape.GetDim(i);
    }
    int32_t axisLength = inputShape.GetDim(i);
    for(i++;i<dimSize;i++){
        batchLength*=inputShape.GetDim(i);
    }

    const auto runtime_attrs=context->GetAttrs();
    bool exclusive=*(runtime_attrs->GetBool(0));
    bool reverse=*(runtime_attrs->GetBool(1));
    
    uint32_t ALIGN_NUM = 512 / typeSize;
    uint32_t totalLengthAligned = (totalLength + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    
    int32_t loopCnt=0;
    uint32_t fullTileLength = 0;
    uint32_t lastTileLength = 0;
    int32_t fullCnt=0,lastCnt=0;
    uint64_t unsignedSize=0;
    int64_t ub_size=0, reservedUB = 512;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, unsignedSize);
    ub_size=(int64_t)unsignedSize;
    ub_size-=reservedUB;

    int64_t wholeBatch = batchLength * axisLength;
    int64_t num = ub_size / 2 / typeSize / batchLength;
    if(num>=3){
        //ok
        uint32_t ub_length = 0;
        ALIGN_NUM = 32 / typeSize;
        if(inputShape.GetDim(axis)==1){
            //don't change it for now
		//key=0;
            ALIGN_NUM = 512 / typeSize;
            ub_length = ( ( (ub_size) / typeSize / UB_SHARING_NUM) / ALIGN_NUM * ALIGN_NUM );
            loopCnt = (totalLength + ub_length - 1) / ub_length;
            //both aligned because of pre-processing
            fullTileLength = ub_length;
            lastTileLength = totalLength - fullTileLength * (loopCnt - 1);
            if(loopCnt==1){
                fullTileLength=0;
            }
        }else if( (batchLength % ALIGN_NUM == 0) && (batchLength * typeSize >= 256) ){
            key+=4;
            //good for UB copying
            ALIGN_NUM = 512 / typeSize;

            reservedUB = (batchLength * typeSize + 511) / 512 * 512;
            ub_size-=reservedUB;

            ub_length = ( ( (ub_size) / typeSize / UB_SHARING_NUM) / ALIGN_NUM * ALIGN_NUM );
            if(ub_length>=wholeBatch){
                //process multiple count of batch
                fullCnt = ub_length / wholeBatch;
                loopCnt = (batchCnt + fullCnt - 1) / fullCnt;
                lastCnt = batchCnt - (loopCnt - 1) * fullCnt;
                if(loopCnt==1){
                    fullCnt = 0;
                }
                fullTileLength = ( fullCnt * wholeBatch + ALIGN_NUM - 1 ) / ALIGN_NUM * ALIGN_NUM;
                lastTileLength = ( lastCnt * wholeBatch + ALIGN_NUM - 1 ) / ALIGN_NUM * ALIGN_NUM;
                
            }else{
                //process a batch along the axis
                fullCnt = ub_length / batchLength;
                loopCnt = (axisLength + fullCnt - 1) / fullCnt;
                lastCnt = axisLength - (loopCnt - 1) * fullCnt;
                if(loopCnt==1){
                    fullCnt = 0;
                }
                fullTileLength = ( fullCnt * batchLength + ALIGN_NUM - 1 ) / ALIGN_NUM * ALIGN_NUM;
                lastTileLength = ( lastCnt * batchLength + ALIGN_NUM - 1 ) / ALIGN_NUM * ALIGN_NUM;
                //need to use an extra local tensor
                lastCnt=-lastCnt;
            }
        }else{
            //calculate simply, no need to configure
        }
    }else{
        //naively, or get two input(x,y) each copyIn
        //key=itself;
    }

    tiling.set_totalLengthAligned(totalLengthAligned);

    tiling.set_loopCnt(loopCnt);
    tiling.set_fullTileLength(fullTileLength);
    tiling.set_lastTileLength(lastTileLength);
    tiling.set_fullCnt(fullCnt);
    tiling.set_lastCnt(lastCnt);
    
    tiling.set_batchCnt(batchCnt);
    tiling.set_axisLength(axisLength);
    tiling.set_batchLength(batchLength);
    tiling.set_exclusive(exclusive);
    tiling.set_reverse(reverse);

    tiling.set_alignNum(ALIGN_NUM);
    tiling.set_typeSize(typeSize);
    
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    context->SetTilingKey(key);
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
/*	
    std::cout<<"Key is "<<key<<std::endl;
    
    std::cout<<"LoopCnt is "<<loopCnt<<std::endl;
    std::cout<<"FullTileLength is "<<fullTileLength<<std::endl;
    std::cout<<"LastTileLength is "<<lastTileLength<<std::endl;
    std::cout<<"FullCnt is "<<fullCnt<<std::endl;
    std::cout<<"LastCnt is "<<lastCnt<<std::endl;
    
    std::cout<<"Axis is "<<axis<<std::endl;
    std::cout<<"Exclusive is "<<exclusive<<std::endl;
    std::cout<<"Reverse is "<<reverse<<std::endl;

    std::cout<<"BatchCnt is "<<batchCnt<<std::endl;
    std::cout<<"AxisLength is "<<axisLength<<std::endl;
    std::cout<<"BatchLength is "<<batchLength<<std::endl;
    std::cout<<"TotalLength is "<<totalLength<<std::endl;
  */  
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Cumsum : public OpDef {
public:
    explicit Cumsum(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("axis")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("exclusive").AttrType(OPTIONAL).Bool(0);
        this->Attr("reverse").AttrType(OPTIONAL).Bool(0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Cumsum);
}
